from .base_exchange import *
from pybit.unified_trading import HTTP, WebSocket
import asyncio
import time

class BybitExchange(BaseExchange):
    def __init__(self, api_key: str = "", api_secret: str = "", is_futures: bool = False, testnet: bool = False):
        super().__init__(api_key, api_secret)
        self.name = "Bybit"
        self.is_futures = is_futures
        self.testnet = testnet
        
        # 初始化REST API客户端
        self.client = HTTP(
            testnet=testnet,
            api_key=api_key,
            api_secret=api_secret
        )
        
        # 初始化WebSocket
        self.ws_public = None
        self.ws_private = None
        
        # 缓存数据
        self.tickers = {}
        self.orderbooks = {}
        self.balances = {}
        self.positions = {}
        
    async def connect_ws(self):
        """连接WebSocket"""
        try:
            # 公共频道WebSocket
            self.ws_public = WebSocket(
                testnet=self.testnet,
                channel_type="linear"
            )
            
            # 私有频道WebSocket
            self.ws_private = WebSocket(
                testnet=self.testnet,
                api_key=self.api_key,
                api_secret=self.api_secret,
                channel_type="private"
            )
            
            # 订阅公共频道
            await self._subscribe_public_channels()
            # 订阅私有频道
            await self._subscribe_private_channels()
            
        except Exception as e:
            raise Exception(f"Failed to connect WebSocket: {str(e)}")
            
    async def _subscribe_public_channels(self):
        """订阅公共频道"""
        try:
            # 订阅行情数据
            self.ws_public.ticker_stream(
                symbol="BTCUSDT",
                callback=self._handle_ticker
            )
            
            # 订阅深度数据
            self.ws_public.orderbook_stream(
                symbol="BTCUSDT",
                callback=self._handle_orderbook
            )
            
        except Exception as e:
            raise Exception(f"Failed to subscribe public channels: {str(e)}")
            
    async def _subscribe_private_channels(self):
        """订阅私有频道"""
        try:
            # 订阅持仓更新
            self.ws_private.position_stream(
                callback=self._handle_position
            )
            
            # 订阅余额更新
            self.ws_private.wallet_stream(
                callback=self._handle_wallet
            )
            
        except Exception as e:
            raise Exception(f"Failed to subscribe private channels: {str(e)}")
            
    def _handle_ticker(self, message):
        """处理行情数据"""
        try:
            data = message['data']
            symbol = data['symbol']
            self.tickers[symbol] = {
                'last': float(data['lastPrice']),
                'bid': float(data['bidPrice']),
                'ask': float(data['askPrice']),
                'volume': float(data['volume24h']),
                'timestamp': data['timestamp']
            }
        except Exception as e:
            print(f"Error handling ticker: {e}")
            
    def _handle_orderbook(self, message):
        """处理深度数据"""
        try:
            data = message['data']
            symbol = data['symbol']
            self.orderbooks[symbol] = {
                'asks': [(float(price), float(qty)) for price, qty in data['asks']],
                'bids': [(float(price), float(qty)) for price, qty in data['bids']],
                'timestamp': data['timestamp']
            }
        except Exception as e:
            print(f"Error handling orderbook: {e}")
            
    def _handle_position(self, message):
        """处理持仓更新"""
        try:
            data = message['data']
            symbol = data['symbol']
            self.positions[symbol] = Position(
                symbol=symbol,
                size=float(data['size']),
                entry_price=float(data['entryPrice']),
                leverage=float(data['leverage']),
                margin=float(data['positionMargin']),
                unrealized_pnl=float(data['unrealisedPnl']),
                liquidation_price=float(data['liqPrice']),
                timestamp=time.time()
            )
        except Exception as e:
            print(f"Error handling position: {e}")
            
    def _handle_wallet(self, message):
        """处理余额更新"""
        try:
            data = message['data']
            asset = data['coin']
            self.balances[asset] = Balance(
                asset=asset,
                free=float(data['availableBalance']),
                locked=float(data['walletBalance']) - float(data['availableBalance']),
                total=float(data['walletBalance']),
                timestamp=time.time()
            )
        except Exception as e:
            print(f"Error handling wallet: {e}")
            
    async def get_ticker(self, symbol: str) -> Ticker:
        """获取最新行情"""
        try:
            response = self.client.get_tickers(
                category="linear" if self.is_futures else "spot",
                symbol=symbol
            )
            
            ticker_data = response['result']['list'][0]
            return Ticker(
                symbol=symbol,
                last_price=float(ticker_data['lastPrice']),
                bid_price=float(ticker_data['bidPrice']),
                ask_price=float(ticker_data['askPrice']),
                volume_24h=float(ticker_data['volume24h']),
                timestamp=time.time()
            )
        except Exception as e:
            raise Exception(f"Failed to get ticker: {str(e)}")
            
    async def get_orderbook(self, symbol: str, limit: int = 20) -> OrderBook:
        """获取订单簿"""
        try:
            response = self.client.get_orderbook(
                category="linear" if self.is_futures else "spot",
                symbol=symbol,
                limit=limit
            )
            
            return OrderBook(
                asks=[(float(price), float(qty)) for price, qty in response['result']['a']],
                bids=[(float(price), float(qty)) for price, qty in response['result']['b']],
                timestamp=time.time()
            )
        except Exception as e:
            raise Exception(f"Failed to get orderbook: {str(e)}")
            
    async def get_balance(self, asset: str = "") -> Dict[str, Balance]:
        """获取账户余额"""
        try:
            response = self.client.get_wallet_balance(
                accountType="UNIFIED",
                coin=asset if asset else None
            )
            
            balances = {}
            for coin in response['result']['list']:
                asset = coin['coin']
                balances[asset] = Balance(
                    asset=asset,
                    free=float(coin['availableToWithdraw']),
                    locked=float(coin['walletBalance']) - float(coin['availableToWithdraw']),
                    total=float(coin['walletBalance']),
                    timestamp=time.time()
                )
            return balances
        except Exception as e:
            raise Exception(f"Failed to get balance: {str(e)}")
            
    async def get_positions(self, symbol: str = "") -> Dict[str, Position]:
        """获取持仓信息"""
        try:
            response = self.client.get_positions(
                category="linear",
                symbol=symbol if symbol else None
            )
            
            positions = {}
            for pos in response['result']['list']:
                symbol = pos['symbol']
                positions[symbol] = Position(
                    symbol=symbol,
                    size=float(pos['size']),
                    entry_price=float(pos['entryPrice']),
                    leverage=float(pos['leverage']),
                    margin=float(pos['positionMargin']),
                    unrealized_pnl=float(pos['unrealisedPnl']),
                    liquidation_price=float(pos['liqPrice']),
                    timestamp=time.time()
                )
            return positions
        except Exception as e:
            raise Exception(f"Failed to get positions: {str(e)}") 